from traits.api import *
from traitsui.api import *
from mayavi.core.api import Engine, PipelineBase
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, SceneEditor
import numpy as np
from numpy import array, random, linspace, pi, ravel, cos, sin, empty, mgrid, arange
from mayavi import mlab
import networkx as nx
from tvtk.api import tvtk
from mayavi.sources.builtin_surface import BuiltinSurface
import csv
from scipy.special import sph_harm


class Form(HasTraits):
    scene = Instance(MlabSceneModel, ())

    buttonClear = Button('清除视景')

    def _buttonClear_fired(self):
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        mlab.clf()

    fun = Enum('atomic_orbital',
               'boy',
               'polydata',
               'polydata1',
               'polydata2',
               'polydata3',
               'image_data',
               'rectilinear_grid',
               'structured_grid',
               'unstructured_grid',
               'delaunay_graph',
               'custom_colormap',
               'fluidArrow_grid',
               'image_cursor_filter',
               'julia_set',
               'lorenz',
               'RCS_3D',
               'select_red_balls',
               'spherical_harmonics',
               'three3d_wigner',
               'threevisualization1',
               'threevisualization2',
               'testline',
               )

    def traits_init(self):
        self._fun_changed()

    def _fun_changed(self):
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        mlab.clf()
        match self.fun:
            case 'atomic_orbital':
                x, y, z = np.ogrid[- .5:.5:200j, - .5:.5:200j, - .5:.5:200j]
                r = np.sqrt(x ** 2 + y ** 2 + z ** 2)
                # Generalized Laguerre polynomial (3, 2)
                L = - r ** 3 / 6 + 5. / 2 * r ** 2 - 10 * r + 6
                # Spherical harmonic (3, 2)
                Y = (x + y * 1j) ** 2 * z / r ** 3
                Phi = L * Y * np.exp(- r) * r ** 2
                # Plot it ####################################################################
                # We create a scalar field with the module of Phi as the scalar
                src = mlab.pipeline.scalar_field(np.abs(Phi))
                src.image_data.point_data.add_array(np.angle(Phi).T.ravel())
                # We need to give a name to our new dataset.
                src.image_data.point_data.get_array(1).name = 'angle'
                # Make sure that the dataset is up to date with the different arrays:
                src.update()
                # We select the 'scalar' attribute, ie the norm of Phi
                src2 = mlab.pipeline.set_active_attribute(src, point_scalars='scalar')
                # Cut isosurfaces of the norm
                contour = mlab.pipeline.contour(src2)
                # Now we select the 'angle' attribute, ie the phase of Phi
                contour2 = mlab.pipeline.set_active_attribute(contour, point_scalars='angle')
                # And we display the surface. The colormap is the current attribute: the phase.
                mlab.pipeline.surface(contour2, colormap='hsv')
                # self.scene.mlab.colorbar(title='Phase', orientation='vertical', nb_labels=3)

            case 'boy':
                u, v = np.mgrid[- 0.035:np.pi:0.01, - 0.05:np.pi:0.01]
                X = 2 / 3. * (np.cos(u) * np.cos(2 * v) + np.sqrt(2) * np.sin(u) * np.cos(v)) * np.cos(u) / (np.sqrt(2) - np.sin(2 * u) * np.sin(3 * v))
                Y = 2 / 3. * (np.cos(u) * np.sin(2 * v) - np.sqrt(2) * np.sin(u) * np.sin(v)) * np.cos(u) / (np.sqrt(2) - np.sin(2 * u) * np.sin(3 * v))
                Z = -np.sqrt(2) * np.cos(u) * np.cos(u) / (np.sqrt(2) - np.sin(2 * u) * np.sin(3 * v))
                S = np.sin(u)
                mlab.mesh(X, Y, Z, scalars=S, colormap='YlGnBu', )

            case 'polydata':
                def polydata():
                    # The numpy array data.
                    points = np.array(
                        [[0, -0.5, 0], [1.5, 0, 0], [0, 1, 0], [0, 0, 0.5], [-1, -1.5, 0.1], [0, -1, 0.5], [-1, -0.5, 0],
                         [1, 0.8, 0]], 'f')
                    triangles = np.array(
                        [[0, 1, 3], [1, 2, 3], [1, 0, 5], [2, 3, 4], [3, 0, 4], [0, 5, 4], [2, 4, 6], [2, 1, 7]])
                    scalars = np.random.random(points.shape)
                    # The TVTK dataset.
                    mesh = tvtk.PolyData(points=points, polys=triangles)
                    mesh.point_data.scalars = scalars
                    mesh.point_data.scalars.name = 'scalars'
                    return mesh

                def view(dataset):
                    surf = mlab.pipeline.surface(dataset, opacity=0.1)
                    mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

                view(polydata())

            case 'polydata1':
                points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
                triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
                scalars = np.array([0, .1, .2, .3, .4], dtype=float)
                # scalars = np.array([[0, 1, 0]])
                mesh = tvtk.PolyData(points=points, polys=triangles)
                mesh.point_data.scalars = scalars
                # mesh.cell_data.scalars = scalars
                surf = mlab.pipeline.surface(mesh)
                mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

            case 'polydata2':
                points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
                triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
                scalars = np.array([0, .1, .2, .3, .4], dtype=float)
                surf = mlab.triangular_mesh(points[:, 0], points[:, 1], points[:, 2], triangles, scalars=scalars)
                mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

            case 'polydata3':
                points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
                triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
                scalars = np.array([[0, .5, 1]])
                mesh = tvtk.PolyData(points=points, polys=triangles)
                mesh.cell_data.scalars = scalars
                surf = mlab.pipeline.surface(mesh)
                mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

            case 'image_data':
                def image_data():
                    data = np.random.random((3, 3, 3))
                    i = tvtk.ImageData(spacing=(1, 1, 1), origin=(0, 0, 0))
                    i.point_data.scalars = data.ravel()
                    i.point_data.scalars.name = 'scalars'
                    i.dimensions = data.shape
                    return i

                def view(dataset):
                    surf = mlab.pipeline.surface(dataset, opacity=0.1)
                    mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

                view(image_data())

            case 'rectilinear_grid':
                def rectilinear_grid():
                    data = np.random.random((3, 3, 3))
                    r = tvtk.RectilinearGrid()
                    r.point_data.scalars = data.ravel()
                    r.point_data.scalars.name = 'scalars'
                    r.dimensions = data.shape
                    r.x_coordinates = np.array((0, 0.7, 1.4))
                    r.y_coordinates = np.array((0, 1, 3))
                    r.z_coordinates = np.array((0, .5, 2))
                    return r

                def view(dataset):
                    surf = mlab.pipeline.surface(dataset, opacity=0.1)
                    mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

                view(rectilinear_grid())

            case 'structured_grid':
                def generate_annulus(r, theta, z):
                    """ Generate points for structured grid for a cylindrical annular
                        volume.  This method is useful for generating a unstructured
                        cylindrical mesh for VTK (and perhaps other tools).
                    """
                    # Find the x values and y values for each plane.
                    x_plane = (np.cos(theta) * r[:, None]).ravel()
                    y_plane = (np.sin(theta) * r[:, None]).ravel()
                    # Allocate an array for all the points.  We'll have len(x_plane)
                    # points on each plane, and we have a plane for each z value, so
                    # we need len(x_plane)*len(z) points.
                    points = np.empty([len(x_plane) * len(z), 3])
                    # Loop through the points for each plane and fill them with the
                    # correct x,y,z values.
                    start = 0
                    for z_plane in z:
                        end = start + len(x_plane)
                        # slice out a plane of the output points and fill it
                        # with the x,y, and z values for this plane.  The x,y
                        # values are the same for every plane.  The z value
                        # is set to the current z
                        plane_points = points[start:end]
                        plane_points[:, 0] = x_plane
                        plane_points[:, 1] = y_plane
                        plane_points[:, 2] = z_plane
                        start = end
                    return points

                def structured_grid():
                    # Make the data.
                    dims = (3, 4, 3)
                    r = np.linspace(5, 15, dims[0])
                    theta = np.linspace(0, 0.5 * np.pi, dims[1])
                    z = np.linspace(0, 10, dims[2])
                    pts = generate_annulus(r, theta, z)
                    sgrid = tvtk.StructuredGrid(dimensions=(dims[1], dims[0], dims[2]))
                    sgrid.points = pts
                    s = np.random.random((dims[0] * dims[1] * dims[2]))
                    sgrid.point_data.scalars = np.ravel(s.copy())
                    sgrid.point_data.scalars.name = 'scalars'
                    return sgrid

                def view(dataset):
                    surf = mlab.pipeline.surface(dataset, opacity=0.1)
                    mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

                view(structured_grid())

            case 'unstructured_grid':
                def generate_annulus(r, theta, z):
                    """ Generate points for structured grid for a cylindrical annular
                        volume.  This method is useful for generating a unstructured
                        cylindrical mesh for VTK (and perhaps other tools).
                    """
                    # Find the x values and y values for each plane.
                    x_plane = (np.cos(theta) * r[:, None]).ravel()
                    y_plane = (np.sin(theta) * r[:, None]).ravel()
                    # Allocate an array for all the points.  We'll have len(x_plane)
                    # points on each plane, and we have a plane for each z value, so
                    # we need len(x_plane)*len(z) points.
                    points = np.empty([len(x_plane) * len(z), 3])
                    # Loop through the points for each plane and fill them with the
                    # correct x,y,z values.
                    start = 0
                    for z_plane in z:
                        end = start + len(x_plane)
                        # slice out a plane of the output points and fill it
                        # with the x,y, and z values for this plane.  The x,y
                        # values are the same for every plane.  The z value
                        # is set to the current z
                        plane_points = points[start:end]
                        plane_points[:, 0] = x_plane
                        plane_points[:, 1] = y_plane
                        plane_points[:, 2] = z_plane
                        start = end
                    return points

                def unstructured_grid():
                    points = np.array([[0, 1.2, 0.6], [1, 0, 0], [0, 1, 0], [1, 1, 1],  # tetra
                                       [1, 0, -0.5], [2, 0, 0], [2, 1.5, 0], [0, 1, 0],
                                       [1, 0, 0], [1.5, -0.2, 1], [1.6, 1, 1.5], [1, 1, 1],  # Hex
                                       ], 'f')
                    # The cells
                    cells = np.array([4, 0, 1, 2, 3,  # tetra
                                      8, 4, 5, 6, 7, 8, 9, 10, 11  # hex
                                      ])
                    # The offsets for the cells, i.e. the indices where the cells
                    # start.
                    offset = np.array([0, 5])
                    tetra_type = tvtk.Tetra().cell_type  # VTK_TETRA == 10
                    hex_type = tvtk.Hexahedron().cell_type  # VTK_HEXAHEDRON == 12
                    cell_types = np.array([tetra_type, hex_type])
                    # Create the array of cells unambiguously.
                    cell_array = tvtk.CellArray()
                    cell_array.set_cells(2, cells)
                    # Now create the UG.
                    ug = tvtk.UnstructuredGrid(points=points)
                    # Now just set the cell types and reuse the ug locations and cells.
                    ug.set_cells(cell_types, offset, cell_array)
                    scalars = np.random.random(points.shape[0])
                    ug.point_data.scalars = scalars
                    ug.point_data.scalars.name = 'scalars'
                    return ug

                def view(dataset):
                    surf = mlab.pipeline.surface(dataset, opacity=0.1)
                    mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

                view(unstructured_grid())

            case 'delaunay_graph':
                def compute_delaunay_edges(x, y, z, visualize=False):
                    """ Given 3-D points, returns the edges of their
                        Delaunay triangulation.
                        Parameters
                        -----------
                        x: ndarray
                            x coordinates of the points
                        y: ndarray
                            y coordinates of the points
                        z: ndarray
                            z coordinates of the points
                        Returns
                        ---------
                        new_x: ndarray
                            new x coordinates of the points (same coords but different
                            assignment of points)
                        new_y: ndarray
                            new y coordinates of the points (same coords but different
                            assignment of points)
                        new_z: ndarray
                            new z coordinates of the points (same coords but different
                            assignment of points)
                        edges: 2D ndarray.
                            The indices of the edges of the Delaunay triangulation as a
                            (N, 2) array [[pair1_index1, pair1_index2],
                                          [pair2_index1, pair2_index2],
                                          ...                         ]
                    """
                    if visualize:
                        vtk_source = mlab.points3d(x, y, z, opacity=0.3, mode='2dvertex')
                        vtk_source.actor.property.point_size = 3
                    else:
                        vtk_source = mlab.pipeline.scalar_scatter(x, y, z, figure=False)
                    delaunay = mlab.pipeline.delaunay3d(vtk_source)
                    delaunay.filter.offset = 999  # seems more reliable than the default
                    edges = mlab.pipeline.extract_edges(delaunay)
                    if visualize:
                        mlab.pipeline.surface(edges, opacity=0.3, line_width=3)
                    # We extract the output array. the 'points' attribute itself
                    # is a TVTK array, that we convert to a numpy array using
                    # its 'to_array' method.
                    dataset = edges.get_output_dataset()
                    new_x, new_y, new_z = dataset.points.to_array().T
                    lines = dataset.lines.to_array()
                    return new_x, new_y, new_z, np.array([lines[1::3], lines[2::3]]).T

                def graph_plot(x, y, z, start_idx, end_idx, edge_scalars=None, **kwargs):
                    """ Show the graph edges using Mayavi
                        Parameters
                        -----------
                        x: ndarray
                            x coordinates of the points
                        y: ndarray
                            y coordinates of the points
                        z: ndarray
                            z coordinates of the points
                        edge_scalars: ndarray, optional
                            optional data to give the color of the edges.
                        kwargs:
                            extra keyword arguments are passed to quiver3d.
                    """
                    vec = mlab.quiver3d(x[start_idx], y[start_idx], z[start_idx],
                                        x[end_idx] - x[start_idx], y[end_idx] - y[start_idx], z[end_idx] - z[start_idx],
                                        scalars=edge_scalars, mode='2ddash', scale_factor=1, **kwargs)
                    if edge_scalars is not None:
                        vec.glyph.color_mode = 'color_by_scalar'
                    return vec

                def build_geometric_graph(x, y, z, edges):
                    """ Build a NetworkX graph with xyz node coordinates and the node indices
                        of the end nodes.
                        Parameters
                        -----------
                        x: ndarray
                            x coordinates of the points
                        y: ndarray
                            y coordinates of the points
                        z: ndarray
                            z coordinates of the points
                        edges: the (2, N) array returned by compute_delaunay_edges()
                            containing node indices of the end nodes. Weights are applied to
                            the edges based on their euclidean length for use by the MST
                            algorithm.
                        Returns
                        ---------
                        g: A NetworkX undirected graph
                        Notes
                        ------
                        We don't bother putting the coordinates into the NX graph.
                        Instead the graph node is an index to the column.
                    """
                    xyz = np.array((x, y, z))

                    def euclidean_dist(i, j):
                        d = xyz[:, i] - xyz[:, j]
                        return np.sqrt(np.dot(d, d))

                    g = nx.Graph()
                    for i, j in edges:
                        if nx.__version__.split('.')[0] > '0':
                            g.add_edge(i, j, weight=euclidean_dist(i, j))
                        else:
                            g.add_edge(i, j, euclidean_dist(i, j))
                    return g

                def points_on_sphere(N):
                    """ Generate N evenly distributed points on the unit sphere centered at
                        the origin. Uses the 'Golden Spiral'.
                        Code by Chris Colbert from the numpy-discussion list.
                    """
                    phi = (1 + np.sqrt(5)) / 2  # the golden ratio
                    long_incr = 2 * np.pi / phi  # how much to increment the longitude
                    dz = 2.0 / float(N)  # a unit sphere has diameter 2
                    bands = np.arange(N)  # each band will have one point placed on it
                    z = bands * dz - 1 + (dz / 2)  # the height z of each band/point
                    r = np.sqrt(1 - z * z)  # project onto xy-plane
                    az = bands * long_incr  # azimuthal angle of point modulo 2 pi
                    x = r * np.cos(az)
                    y = r * np.sin(az)
                    return x, y, z

                    # ################################################################################
                    # if __name__ == '__main__':
                    # generate some points

                x, y, z = points_on_sphere(50)
                # Avoid triangulation problems on the sphere
                z *= 1.01
                # mlab.figure(1, bgcolor=(0, 0, 0))
                # mlab.clf()
                # Now get the Delaunay Triangulation from vtk via mayavi mlab. Vtk stores
                # its points in a different order so overwrite ours to match the edges
                new_x, new_y, new_z, edges = compute_delaunay_edges(x, y, z, visualize=True)
                assert (x.shape == new_x.shape)  # check triangulation got everything
                x, y, z = new_x, new_y, new_z
                if nx.__version__ < '0.99':
                    raise ImportError('The version of NetworkX must be at least '
                                      '0.99 to run this example')
                    # Make a NetworkX graph out of our point and edge data
                g = build_geometric_graph(x, y, z, edges)
                # Compute minimum spanning tree using networkx
                # nx.mst returns an edge generator
                edges = nx.minimum_spanning_tree(g).edges(data=True)
                start_idx, end_idx, _ = np.array(list(edges)).T
                start_idx = start_idx.astype(int)
                end_idx = end_idx.astype(int)
                # Plot this with Mayavi
                graph_plot(x, y, z, start_idx, end_idx, edge_scalars=z[start_idx], opacity=0.8, colormap='summer',
                           line_width=4, )

            case 'custom_colormap':
                x, y = np.mgrid[-10:10:200j, -10:10:200j]
                z = 100 * np.sin(x * y) / (x * y)
                # self.scene.mlab.figure(bgcolor=(1, 1, 1))
                surf = mlab.surf(z, colormap='cool')
                lut = surf.module_manager.scalar_lut_manager.lut.table.to_array()
                lut[:, -1] = np.linspace(0, 255, 256)
                surf.module_manager.scalar_lut_manager.lut.table = lut

            case 'fluidArrow_grid':
                x, y, z = np.mgrid[1:6:11j, 0:4:13j, 0:3:6j]
                base = x[..., 0] + y[..., 0]
                for i in range(z.shape[2]):
                    z[..., i] = base * 0.25 * i
                pts = np.empty(z.shape + (3,), dtype=float)
                pts[..., 0] = x
                pts[..., 1] = y
                pts[..., 2] = z
                scalars = x * x + y * y + z * z
                vectors = np.empty(z.shape + (3,), dtype=float)
                vectors[..., 0] = (4 - y * 2)
                vectors[..., 1] = (x * 3 - 12)
                vectors[..., 2] = np.sin(z * np.pi)
                pts = pts.transpose(2, 1, 0, 3).copy()
                pts.shape = pts.size // 3, 3
                scalars = scalars.T.copy()
                vectors = vectors.transpose(2, 1, 0, 3).copy()
                vectors.shape = vectors.size // 3, 3
                sg = tvtk.StructuredGrid(dimensions=x.shape, points=pts)
                sg.point_data.scalars = scalars.ravel()
                sg.point_data.scalars.name = 'temperature'
                sg.point_data.vectors = vectors
                sg.point_data.vectors.name = 'velocity'
                d = mlab.pipeline.add_dataset(sg)
                gx = mlab.pipeline.grid_plane(d)
                gy = mlab.pipeline.grid_plane(d)
                gy.grid_plane.axis = 'y'
                gz = mlab.pipeline.grid_plane(d)
                gz.grid_plane.axis = 'z'
                iso = mlab.pipeline.iso_surface(d)
                iso.contour.maximum_contour = 75.0
                vec = mlab.pipeline.vectors(d)
                vec.glyph.mask_input_points = True
                vec.glyph.glyph.scale_factor = 1.5

            case 'image_cursor_filter':
                x, y, z = np.ogrid[0:1:40j, 0:1:40j, 0:1:40j]
                sig = 0.5
                center = 0.5
                g = np.exp(-((x - center) ** 2 + (y - center) ** 2 + (z - center) ** 2) / (2 * sig ** 2))
                s = mlab.pipeline.scalar_field(g)
                cursor = mlab.pipeline.user_defined(s, filter='ImageCursor3D')
                cursor.filter.cursor_position = np.array([20, 20, 20])
                cursor.filter.cursor_value = 0
                cursor.filter.cursor_radius = 10
                ipw = mlab.pipeline.image_plane_widget(cursor, plane_orientation='x_axes', slice_index=20)
                mlab.colorbar()
                mlab.view(15, 70, 100, [20, 20, 20])

            case 'julia_set':
                x, y = np.ogrid[-1.5:0.5:500j, -1:1:500j]
                z = x + 1j * y
                julia = np.zeros(z.shape)
                for i in range(50):
                    z = z ** 2 - 0.70176 - 0.3842j
                    julia += 1 / float(2 + i) * (z * np.conj(z) > 4)
                mlab.surf(julia, colormap='gist_earth', warp_scale='auto', vmax=1.5)

            case 'lorenz':
                def lorenz(x, y, z, s=10., r=28., b=8. / 3.):
                    """The Lorenz system."""
                    u = s * (y - x)
                    v = r * x - y - x * z
                    w = x * y - b * z
                    return u, v, w

                x, y, z = np.mgrid[-50:50:100j, -50:50:100j, -10:60:70j]
                u, v, w = lorenz(x, y, z)
                # fig = mlab.figure(size=(400, 300), bgcolor=(0, 0, 0))
                f = mlab.flow(x, y, z, u, v, w, line_width=3, colormap='Paired')
                f.module_manager.scalar_lut_manager.reverse_lut = True
                f.stream_tracer.integration_direction = 'both'
                # f.stream_tracer.maximum_propagation = 200
                src = f.mlab_source.m_data
                e = mlab.pipeline.extract_vector_components(src)
                e.component = 'z-component'
                zc = mlab.pipeline.iso_surface(e, opacity=0.5, contours=[0, ], color=(0.6, 1, 0.2))
                zc.actor.property.backface_culling = True

            case 'RCS_3D':
                N = 300
                t = np.linspace(-2 * np.pi, 2 * np.pi, N)
                x = list()
                y = list()
                z = list()
                s = list()
                connections = list()
                index = 0
                for i in range(50):
                    x.append(np.sin(t))
                    y.append(np.cos((2 + .02 * i) * t))
                    z.append(np.cos((3 + .02 * i) * t))
                    s.append(t)
                    connections.append(np.vstack(
                        [np.arange(index, index + N - 1.5), np.arange(index + 1, index + N - .5)]).T)
                    index += N
                x = np.hstack(x)
                y = np.hstack(y)
                z = np.hstack(z)
                s = np.hstack(s)
                connections = np.vstack(connections)
                src = mlab.pipeline.scalar_scatter(x, y, z, s)
                src.mlab_source.dataset.lines = connections
                src.update()
                lines = mlab.pipeline.stripper(src)
                mlab.pipeline.surface(lines, colormap='Accent', line_width=1, opacity=.4)

            case 'select_red_balls':
                figure = mlab.gcf()
                # mlab.clf()
                # figure.scene.disable_render = True
                x1, y1, z1 = np.random.random((3, 10))
                red_glyphs = mlab.points3d(x1, y1, z1, color=(1, 0, 0), resolution=20)
                x2, y2, z2 = np.random.random((3, 10))
                white_glyphs = mlab.points3d(x2, y2, z2, color=(0.9, 0.9, 0.9), resolution=20)
                outline = mlab.outline(line_width=3)
                outline.outline_mode = 'cornered'
                outline.bounds = (x1[0] - 0.1, x1[0] + 0.1, y1[0] - 0.1, y1[0] + 0.1, z1[0] - 0.1, z1[0] + 0.1)
                # figure.scene.disable_render = False
                glyph_points = red_glyphs.glyph.glyph_source.glyph_source.output.points.to_array()

                def picker_callback(picker):
                    if picker.actor in red_glyphs.actor.actors:
                        point_id = picker.point_id // glyph_points.shape[0]
                        if point_id != -1:
                            x, y, z = x1[point_id], y1[point_id], z1[point_id]
                            outline.bounds = (x - 0.1, x + 0.1, y - 0.1, y + 0.1, z - 0.1, z + 0.1)

                picker = figure.on_mouse_pick(picker_callback)
                picker.tolerance = 0.01
                mlab.title('Click on red balls')

            case 'spherical_harmonics':
                r = 0.3
                phi, theta = np.mgrid[0:np.pi:101j, 0:2 * np.pi:101j]
                x = r * np.sin(phi) * np.cos(theta)
                y = r * np.sin(phi) * np.sin(theta)
                z = r * np.cos(phi)
                # Represent spherical harmonics on the surface of the sphere
                for n in range(0, 5):
                    for m in range(n):
                        s = sph_harm(m, n, theta, phi).real
                        mlab.mesh(x + m, y + n, z, scalars=s, colormap='jet')
                        s[s < 0] *= 0.97
                        s /= s.max()
                        mlab.mesh(s * x + m, s * y + n, s * z + 1.3, scalars=s, colormap='jet')

            case 'three3d_wigner':
                def cat(x, y, alpha, eta=1, purity=1):
                    return (1 + eta * (
                            np.exp(-x ** 2 - (y - alpha) ** 2) + np.exp(-x ** 2 - (y + alpha) ** 2) + 2 * purity * np.exp(
                        -x ** 2 - y ** 2) * np.cos(2 * alpha * x)) / (2 * (1 + np.exp(- alpha ** 2)))) / 2

                x, y = np.mgrid[-4:4.15:0.1, -4:4.15:0.1]
                cat1 = cat(x, y, 1)
                cat2 = cat(x, y, 2)
                cat3 = cat(x, y, 3)
                cat1 -= 0.5
                cat2 -= 0.5
                cat3 -= 0.5
                cat1_extent = (-14, -6, -4, 4, 0, 5)
                surf_cat1 = mlab.surf(x - 10, y, cat1, colormap='Spectral', warp_scale=5, extent=cat1_extent, vmin=-0.5,
                                      vmax=0.5)
                # mlab.outline(surf_cat1, color=(.7, .7, .7), extent=cat1_extent)
                # mlab.axes(surf_cat1, color=(.7, .7, .7), extent=cat1_extent,
                # ranges=(0, 1, 0, 1, 0, 1), xlabel='', ylabel='',
                # zlabel='Probability',
                # x_axis_visibility=False, z_axis_visibility=False)
                # mlab.text(-18, -4, '1 photon', z=-4, width=0.13)
                cat2_extent = (-4, 4, -4, 4, 0, 5)
                surf_cat2 = mlab.surf(x, y, cat2, colormap='Spectral', warp_scale=5, extent=cat2_extent, vmin=-0.5, vmax=0.5)
                # self.scene.mlab.outline(surf_cat2, color=(0.7, .7, .7), extent=cat2_extent)
                # self.scene.mlab.text(-4, -3, '2 photons', z=-4, width=0.14)
                cat3_extent = (6, 14, -4, 4, 0, 5)
                surf_cat3 = mlab.surf(x + 10, y, cat3, colormap='Spectral', warp_scale=5, extent=cat3_extent, vmin=-0.5,
                                      vmax=0.5)
                # self.scene.mlab.outline(surf_cat3, color=(.7, .7, .7), extent=cat3_extent)
                # self.scene.mlab.text(6, -2.5, '3 photons', z=-4, width=0.14)
                # self.scene.mlab.title('Multi-photons cats Wigner function')

            case 'threevisualization1':
                mlab.test_points3d()

            case 'threevisualization2':
                np.random.seed(12345)
                x = 4 * (np.random.random(500) - 0.5)
                y = 4 * (np.random.random(500) - 0.5)

                def f(x, y):
                    return np.exp(-(x ** 2 + y ** 2))

                z = f(x, y)
                pts = mlab.points3d(x, y, z, z, scale_mode='none', scale_factor=0.2)
                mesh = mlab.pipeline.delaunay2d(pts)
                surf = mlab.pipeline.surface(mesh)
            case 'testline':
                t=np.arange(1000)*np.pi/180
                x = np.cos(t*4)
                y = np.sin(t*4)
                z = t/4
                mlab.plot3d(x, y, z)

    view = View(
        Group(
            Tabbed(
                UItem('scene', editor=SceneEditor(scene_class=MayaviScene)),
            ),
            HGroup(
                UItem('fun'),
                UItem('buttonClear'),
                spring
            ),
        ),
        height=800,
        width=1000,
        resizable=True  # We need this to resize with the parent widget
    )


if __name__ == '__main__':
    Form().configure_traits()
